package com.whz.governance.redislimiter.filter;

import com.beust.jcommander.internal.Lists;
import com.whz.generic.utils.IpUtils;
import com.whz.generic.utils.PathUtils;
import com.whz.governance.redislimiter.command.LeakyBucketResult;
import com.whz.governance.redislimiter.command.TokenBucketResult;
import com.whz.governance.redislimiter.config.RedisRateLimiterProperties.BucketType;
import com.whz.governance.redislimiter.config.RedisRateLimiterProperties.RateLimiterConfig;
import com.whz.governance.redislimiter.tool.LeakyBucketUtils;
import com.whz.governance.redislimiter.tool.TokenBucketUtils;
import lombok.extern.slf4j.Slf4j;
import org.apache.commons.lang.BooleanUtils;
import org.springframework.util.CollectionUtils;
import org.springframework.util.StringUtils;

import javax.servlet.http.HttpServletRequest;
import java.util.Arrays;
import java.util.List;
import java.util.Objects;
import java.util.regex.Pattern;


/**
 * @author whz
 */
@Slf4j
public class RedisRateLimitSupporter implements RateLimitSupporter {

    private static final Pattern COMPILE = Pattern.compile("/");
    public static final String SUFFIX = ":";
    private final List<String> pattern;
    private final RateLimiterConfig config;
    private final List<String> whiteIp;

    public RedisRateLimitSupporter(String pattern, RateLimiterConfig config, List<String> whiteIp) {
        this.config = config;
        this.pattern = split(pattern);
        this.whiteIp = Lists.newArrayList(whiteIp);
    }

    @Override
    public boolean match(String path) {
        if (CollectionUtils.isEmpty(pattern)) {
            return false;
        }
        return pattern.stream().anyMatch(p -> PathUtils.match(p, path));
    }

    @Override
    public boolean allowAccess(final HttpServletRequest request) {

        /*
         * we cannot get remote ip,allow it
         */
        String remoteIp = IpUtils.getRemoteIpAddrByRequest(request);
        if (StringUtils.isEmpty(remoteIp) || isWhiteIps(remoteIp)) {
            log.info("remoteIP: '{}' in white List, allowed access.", remoteIp);
            return true;
        }

        /*
         * if request from inner net, allow it
         */
        if (config.isIntranetEnable() && IpUtils.isInnerIp(remoteIp)) {
            log.info("remoteIP:{} is inner IP.", remoteIp);
            return true;
        }

        /*
         * leaky or token bucket limiter
         */
        String requestUri = request.getRequestURI();
        remoteIp = StringUtils.replace(remoteIp, SUFFIX, "-");
        if (Objects.requireNonNull(config.getAlgorithm()) == BucketType.LEAKY) {
            return BooleanUtils.isTrue(doLeakyAccess(remoteIp, requestUri));
        }
        return BooleanUtils.isTrue(doTokenAccess(remoteIp, requestUri));
    }

    private boolean isWhiteIps(String remoteIp) {
        return whiteIp != null && whiteIp.stream().anyMatch(whiteIp -> whiteIp.equalsIgnoreCase(remoteIp));
    }

    private Boolean doTokenAccess(String remoteIp, String requestUri) {
        try {
            TokenBucketResult result =
                    TokenBucketUtils.execute(getKey(remoteIp, requestUri),
                                             config.getReplenishRate(),
                                             config.getBurstCapacity(),
                                             config.getRequestedTokens(),
                                             config.getChronoUnit());
            log.info("token bucket return:{}", result);
            return result.isPermit();
        } catch (Throwable exception) {
            log.error("token bucket limit remote: " + remoteIp, exception);
        }
        return true;
    }

    private Boolean doLeakyAccess(String remoteIp, String requestUri) {
        try {
            LeakyBucketResult result = LeakyBucketUtils.execute(getKey(remoteIp, requestUri),
                                                                config.getReplenishRate(),
                                                                config.getBurstCapacity(),
                                                                config.getRequestedTokens(),
                                                                config.getChronoUnit());
            log.info("leaky bucket return:{}", result);
            return result.isPermit();
        } catch (Throwable exception) {
            log.error("leaky bucket limit remote: " + remoteIp, exception);
        }
        return true;
    }

    private static List<String> split(String input) {
        return Arrays.asList(StringUtils.tokenizeToStringArray(input, ","));
    }

    private static String getKey(String remoteIp, String requestUri) {
        // 每个接口都有一个桶,独立限流
        String path = COMPILE.matcher(requestUri).replaceAll(SUFFIX);
        if (!path.endsWith(SUFFIX)) {
            path += SUFFIX;
        }
        return path + remoteIp;
    }
}